package cs429_project.clustering;

import java.io.File;
import java.io.IOException;
import java.io.PrintStream;

import java.math.BigDecimal;
import java.math.RoundingMode;

import java.util.Iterator;
import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;

import org.apache.lucene.analysis.standard.StandardAnalyzer;

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.TermFreqVector;

import org.apache.lucene.queryParser.ParseException;
import org.apache.lucene.queryParser.QueryParser;

import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TopDocs;

import org.apache.lucene.store.FSDirectory;

import org.apache.lucene.util.Version;

/**
 * Class that represents a document as a vector mapping its terms to their
 * frequencies.  This class was originally called {@code DocVec2}, but the
 * original {@code DocVec} class has been removed.
 * @author Peter Ferrans
 */
public class DocVec extends TreeMap<String, BigDecimal> implements
        Comparable<DocVec> {

    /** Default precision used on mathematical operations involving
     * {@link java.math.BigDecimal} data types. */
    public static final int DEFAULT_PRECISION = 10;
    /** ID number for the document represented by this vector. */
    public int docId;

    /**
     * Default constructor.  Creates an empty document vector with ID 0.
     */
    public DocVec() {
        super();
        this.docId = 0;
    }

    /**
     * Given the name of a Lucene index, a document ID number, and a field,
     * finds the document with the specified ID from the index, retrieves the
     * text in the field specified, and creates a document vector using the
     * words in that field.  The vector's ID number is also set to be the ID
     * passed in through the parameter.
     * @param indexName
     *    Name of the index to search to find the document to convert to a
     *    document vector.
     * @param docId
     *    ID number for the document to find and for this vector.
     * @param field
     *    Field in the specified Lucene document from which to retrieve text.
     * @throws IOException
     */
    public DocVec(String indexName, int docId, String field)
            throws IOException {
        super();
        this.docId = docId;
        TermFreqVector tfv = getTermFreqVector(indexName, docId, field);
        buildMap(tfv);
    }

    /**
     * Given an array of tokens and an ID number, creates a document vector
     * with the ID number and the tokens as its terms.  The number of times a
     * term appears in the array will be the frequency stored for the term in
     * the document vector.
     * @param tokens
     *    Tokens to be used as terms in this document vector.
     * @param docId
     *    ID number for this document vector.
     */
    public DocVec(String[] tokens, int docId) {
        super();
        this.docId = docId;
        for(int i = 0; i < tokens.length; i++) {
            if(this.containsKey(tokens[i])) {
                BigDecimal newFreq = this.get(tokens[i]).add(BigDecimal.ONE);
                this.put(tokens[i], newFreq);
            } else {
                this.put(tokens[i], BigDecimal.ONE);
            }
        }
    }

    /**
     * Adds a term to the document vector with the specified frequency.  If the
     * term already exists, the frequency here is set to be its new frequency.
     * Functions exactly like {@link #put(java.lang.Object, java.lang.Object)},
     * but the frequency is a {@code double} rather than a
     * {@link java.math.BigDecimal}.
     * @param term
     *    Term to be put into the map.
     * @param frequency
     *    Frequency for the new term.
     * @return
     *    The new frequency for the term put onto the map.
     */
    public BigDecimal put(String term, double frequency) {
        BigDecimal bdFreq = BigDecimal.valueOf(frequency);
        return this.put(term, bdFreq);
    }

    /**
     * Adds entries to this document vector for each term passed in through
     * the parameter.  For each term to be added, if the term does not already
     * exist in the map, it is added with a frequency of 0.  If it does already
     * exist, its frequency is not changed from its current value.
     * @param newTerms
     *    New terms to be added to the document vector for expansion.
     * @return
     *    The number of terms passed in that were not in the document vector at
     *    the start of this method.
     */
    public int expand(Set<String> newTerms) {
        int oldSize = this.size();
        Iterator<String> termIt = newTerms.iterator();
        while(termIt.hasNext()) {
            String term = termIt.next().toLowerCase();
            if(!(this.containsTerm(term))) {
                this.put(term, BigDecimal.ZERO);
            }
        }
        int newSize = this.size();

        return (newSize - oldSize);
    }

    /**
     * Computes the total number of words in this document vector.  It is a sum
     * of all the term frequencies.
     * @return
     *    Number of words appearing in this document vector.
     */
    public BigDecimal numTotalWords() {
        BigDecimal wordCount = BigDecimal.ZERO;
        Iterator<BigDecimal> freqIt = this.values().iterator();
        while(freqIt.hasNext()) {
            BigDecimal freq = freqIt.next();
            wordCount = wordCount.add(freq);
        }

        return wordCount;
    }

    /**
     * Returns a set of all terms in this document vector, including those with
     * zero frequency.
     * @return
     *    Set of all terms in this document vector.
     */
    public Set<String> getTerms() {
        return this.keySet();
    }

    /**
     * Returns a set of all the frequencies in this document vector.
     * @return
     *    Set of all frequencies in this document vector.
     */
    public Set<BigDecimal> getFrequencies() {
        return new TreeSet<BigDecimal>(this.values());
    }

    /**
     * Indicates whether or not this document vector contains a term.  The
     * vector contains a term if it has an entry for it, even if the frequency
     * for that term is zero.
     * @param term
     *    Term for which to check the vector.
     * @return
     *    Whether or not the vector contains the term.
     */
    public boolean containsTerm(String term) {
        return this.containsKey(term);
    }

    /**
     * Indicates whether or not this document vector contains a specified term
     * with a <em>non-zero</em> frequency.
     * @param term
     *    Term for which to check the vector.
     * @return
     *    Whether or not the vector contains the term with a non-zero
     *    frequency.
     */
    public boolean containsNonZeroFreqTerm(String term) {
        boolean contains = this.containsKey(term);
        if(contains) {
            contains = ((this.get(term).compareTo(BigDecimal.ZERO)) != 0);
        }

        return contains;
    }

    /**
     * Prints the data for this vector to the specified print stream.
     * @param stream
     *    Stream to which to print this vector's data.
     */
    public void print(PrintStream stream) {
        stream.println("DOCID: " + this.docId);
        stream.println("TERMS AND FREQUENCIES:");
        String[] terms = this.getTerms().toArray(new String[0]);
        for(int i = 0; i < terms.length; i++) {
            stream.println("\t" + terms[i] + ": " + this.get(terms[i]));
        }
    }

    /**
     * Compares this document vector to another document vector.  Vectors are
     * ranked by document IDs: the greater vector is the one with the greater
     * ID number.
     * @param other
     *    Document vector for comparison.
     * @return
     *    0 if these two document vectors are equal, less than 0 if this vector
     *    less than the other vector, and greater than 0 if this vector is
     *    greater than the other vector.
     */
    public int compareTo(DocVec other) {
        Integer thisDocId = Integer.valueOf(this.docId);
        Integer otherDocId = Integer.valueOf(other.docId);
        
        return thisDocId.compareTo(otherDocId);
    }

    /**
     * Computes the cosine similarity between this document vector and another
     * map of terms to frequencies.  This is used primarily for comparing
     * documents with centroids of clusters.
     * @param other
     *    Map of terms and frequencies for comparison.
     * @return
     *    The cosine similarity of this vector and the other mapping.
     */
    public BigDecimal cosineSimilarity(TreeMap<String, BigDecimal> other) {
        BigDecimal similarity;
        if(other == null) {
            similarity = BigDecimal.valueOf(-1);
        } else {
            BigDecimal dp = this.dotProduct(other);
            BigDecimal size = vectorLength(this).multiply(vectorLength(other));
            similarity = dp.divide(size, DEFAULT_PRECISION, RoundingMode.UP);
        }

        return similarity;
    }

    /**
     * Returns this document vector with all of its frequencies weighted using
     * tf-idf weighting.
     * @param idfs
     *    The idfs (inverse document frequencies) for all the terms in this
     *    document vector.
     * @return
     *    This document vector, with tf-idf weighted frequencies.
     */
    public DocVec tfIdfWeighted(TreeMap<String, BigDecimal> idfs) {
        DocVec tfDocVec = new DocVec();
        Iterator<String> termIt = this.getTerms().iterator();
        tfDocVec.docId = this.docId;
        while(termIt.hasNext()) {
            String term = termIt.next();
            BigDecimal tfIdf = this.get(term).multiply(idfs.get(term));
            tfDocVec.put(term, tfIdf);
        }

        return tfDocVec;
    }

    /**
     * Searches an index with a query and returns as {@link DocVec}s all of the
     * resulting documents.  All of the document vectors that are returned have
     * the exact same set of terms, so they can be used as is in k-means
     * clustering.
     * @param query
     *    Query with which to search the index.
     * @param indexName
     *    Name of Lucene index to search.
     * @param field
     *    Field of documents in index to use for terms.
     * @param maxResults
     *    Maximum number of results to return from the query.
     * @return
     *    A set of the top results from executing a query, in {@link DocVec}
     *    form.
     * @throws IOException
     *    Thrown if there is a problem accessing the index.
     * @throws ParseException
     *    Thrown if there is an error in the query passed in.
     */
    public static TreeSet<DocVec> topQueryResults(String query,
            String indexName, String field, int maxResults) throws IOException,
            ParseException {
        TreeSet<DocVec> results = new TreeSet<DocVec>();

        File indexFile = new File(indexName);
        FSDirectory indexDir = FSDirectory.open(indexFile);
        IndexSearcher searcher = new IndexSearcher(indexDir);

        StandardAnalyzer analyzer = new StandardAnalyzer(Version.LUCENE_30);
        QueryParser parser = new QueryParser(Version.LUCENE_30, field,
                analyzer);
        Query theQuery = parser.parse(query);
        TopDocs topDocs = searcher.search(theQuery, maxResults);

        searcher.close();
        indexDir.close();

        for(int i = 0; i < topDocs.scoreDocs.length; i++) {
            DocVec dv = new DocVec(indexName, topDocs.scoreDocs[i].doc, field);
            results.add(dv);
        }
        results = expandDocVecs(results);
        
        return results;
    }

    /**
     * Expands a set of document vectors so that each vector has a key entry
     * for any term that occurs in any document vector in the set.
     * @param origDocs
     *    Set of document vectors to be expanded.
     * @return
     *    Expanded set of document vectors.
     */
    public static TreeSet<DocVec> expandDocVecs(Set<DocVec> origDocs) {
        TreeSet<String> termSet = new TreeSet<String>();
        Iterator<DocVec> odvIt = origDocs.iterator();
        while(odvIt.hasNext()) {
            DocVec dv = odvIt.next();
            termSet.addAll(dv.getTerms());
        }

        TreeSet<DocVec> expDocs = new TreeSet<DocVec>();
        odvIt = origDocs.iterator();
        while(odvIt.hasNext()) {
            DocVec dv = odvIt.next();
            dv.expand(termSet);
            expDocs.add(dv);
        }

        return expDocs;
    }

    /**
     * Builds the mapping for this document vector using a
     * {@link org.apache.lucene.index.TermFreqVector} of terms and their
     * frequencies, obtained from a Lucene index.
     * @param tfv
     *    {@link org.apache.lucene.index.TermFreqVector} containing the terms
     *    and frequencies that will be used for this document vector.
     */
    private void buildMap(TermFreqVector tfv) {
        String[] terms = tfv.getTerms();
        int[] frequencies = tfv.getTermFrequencies();
        for(int i = 0; i < terms.length; i++) {
            String term = terms[i].toLowerCase();
            BigDecimal frequency = BigDecimal.valueOf(frequencies[i]);
            this.put(term, frequency);
        }
    }

    /**
     * Computes the dot product of this document vector and another mapping of
     * terms and frequencies.
     * @param other
     *    Another mapping of terms and frequencies.  The set of terms used in
     *    the current document vector and {@code other} must be the same.
     * @return
     *    The dot product of this document vector and the other mapping.
     */
    private BigDecimal dotProduct(TreeMap<String, BigDecimal> other) {
        BigDecimal product = BigDecimal.ZERO;
        Iterator<String> termIt = this.getTerms().iterator();
        while(termIt.hasNext()) {
            String term = termIt.next();
            BigDecimal freq1 = this.get(term);
            BigDecimal freq2 = other.get(term);
            if(freq1 == null) {
                freq1 = BigDecimal.ZERO;
            }
            if(freq2 == null) {
                freq2 = BigDecimal.ZERO;
            }
            product = product.add(freq1.multiply(freq2));
        }

        return product;
    }

    /**
     * Given a vector mapping terms to frequencies, calculates the length of
     * the vector.  This length is the square root of the sum of the squares of
     * term frequencies.
     * @param vector
     *    Vector for which we want to compute the length.
     * @return
     *    Length of the vector.
     */
    private static BigDecimal vectorLength(TreeMap<String, BigDecimal> vector) {
        BigDecimal length = BigDecimal.ZERO;
        Iterator<BigDecimal> valIt = vector.values().iterator();
        while(valIt.hasNext()) {
            BigDecimal value = valIt.next();
            value = value.multiply(value);
            length = length.add(value);
        }

        double dblLen = length.doubleValue();
        dblLen = Math.sqrt(dblLen);
        length = BigDecimal.valueOf(dblLen);

        return length;
    }

    /**
     * Given the name of a Lucene index, a document ID number, and the name of
     * a field in that document, fetches the
     * {@link org.apache.lucene.index.TermFreqVector} for that field in the
     * document.
     * @param indexName
     *    Name of index containing the document.
     * @param docId
     *    ID number of the document.
     * @param field
     *    Field in the document.
     * @return
     *    {@link org.apache.lucene.index.TermFreqVector} for the specified
     *    field in the document.
     * @throws IOException
     */
    private static TermFreqVector getTermFreqVector(String indexName, 
            int docId,  String field) throws IOException {
        File indexFile = new File(indexName);
        FSDirectory indexDir = FSDirectory.open(indexFile);
        IndexReader reader = IndexReader.open(indexDir);
        TermFreqVector tfv = reader.getTermFreqVector(docId, field);
        reader.close();
        indexDir.close();
        
        return tfv;
    }
}
